import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import mannwhitneyu, ks_2samp
import os
from config.config import *
from utils.utils import *
from utils.data_split import DataSplitUtil
# ======================================
# 你的数据
df = pd.read_csv(signature_score_csv_path)
train_signature_df,test_signature_df = DataSplitUtil(split_random_state_list[0]).get_train_test_df(df)
signature_X = test_signature_df.drop(columns=exclude_columns)
signature_y = test_signature_df[target_column]
# ======================================
import matplotlib.patches as mpatches

# ======================================
# 可视化和统计检验
results = []
color_palette = sns.color_palette('Set2')

for col in signature_X.columns:
    # # ================== KDE 图
    kde_save_path = opj(signature_analysis_result_path, 'kde')
    md(kde_save_path)
    plt.figure(figsize=(6,4))

    # 分别绘制两组 KDE
    sns.kdeplot(
        data=test_signature_df[test_signature_df['ALN status']==0],
        x=col,
        fill=True,
        color='#66B2FF',
        alpha=0.6,
        common_norm=True,  # 保持与 hue 一致
        bw_adjust=1.0  
    )
    sns.kdeplot(
        data=test_signature_df[test_signature_df['ALN status']==1],
        x=col,
        fill=True,
        color='#FF6666',
        alpha=0.6,
        common_norm=True,  # 保持与 hue 一致
        bw_adjust=1.0  
    )

    # 手动创建图例
    handles = [
        mpatches.Patch(color='#66B2FF', label='Non-ALNM'),
        mpatches.Patch(color='#FF6666', label='ALNM')
    ]
    plt.legend(handles=handles, frameon=False)

    plt.title(f'KDE of {col} Signature')
    plt.tight_layout()
    plt.xlabel(f'{col} Signature Value')
    plt.savefig(os.path.join(kde_save_path, f'{col}.pdf'))
    plt.close()

    # ================== 箱线图
    plt.figure(figsize=(3.5, 4))
    box_save_path = opj(signature_analysis_result_path,'box')
    md(box_save_path)

    # 计算两组 ALN 状态的 p 值（Mann-Whitney U）
    x0 = test_signature_df[test_signature_df['ALN status'] == 0][col]
    x1 = test_signature_df[test_signature_df['ALN status'] == 1][col]
    mw_stat, mw_p = mannwhitneyu(x0, x1, alternative='two-sided')

    x_positions = [0.3, 0.6]  # ALN Negative, ALN Positive
    width = 0.2

    import numpy as np
    color_map = {0: color_palette[0], 1: color_palette[1]}  # 0=ALN-, 1=ALN+
    for i, aln_status in enumerate([0, 1]):
        sub_df = test_signature_df[test_signature_df['ALN status'] == aln_status]

        # 绘制箱体
        plt.boxplot(sub_df[col],
                    positions=[x_positions[i]],
                    widths=width,
                    patch_artist=True,
                     boxprops=dict(facecolor=color_map[aln_status], alpha=0.6, linewidth=0.5),  # 箱体边框粗细
                    medianprops=dict(color='black', linewidth=0.5),  # 中位数线粗细
                    whiskerprops=dict(linewidth=0.5),  # 须的粗细
                    capprops=dict(linewidth=0.5),      # 顶部和底部横线（caps）粗细
                    showfliers=False
                )
        
        # 绘制散点
        x_jitter = np.random.normal(loc=0, scale=0.02, size=len(sub_df))
        plt.scatter(np.full(len(sub_df), x_positions[i]) + x_jitter,
                    sub_df[col],
                    color=color_map[aln_status],
                    alpha=0.7,
                    s=30)


    # 设置标题和坐标
    plt.title(f'{col} Signature Distribution')
    plt.xlim(0, 0.9)
    plt.xticks(x_positions, ['Non-ALNM', 'ALNM'])
    plt.ylabel(f'{col} Signature ')

    # 显著性标注
    if mw_p < 0.001:
        signif = '***'
    elif mw_p < 0.01:
        signif = '**'
    elif mw_p < 0.05:
        signif = '*'
    else:
        signif = ''


    # 横线位置和 p 值
    line_y = 3.4 * 0.75  # 横线高度
    plt.plot(x_positions, [line_y, line_y], color='black', lw=1)  # 横线
    plt.text(0.45, line_y * 1.12, f"p = {format_pval_scientific(mw_p)}",  ha='center', va='bottom', fontsize=11)
    plt.text(0.45, line_y * 1.01, f'{signif}', ha='center', va='bottom', fontsize=9)

    plt.ylim(-3, 3.4)
    # 两端垂直短线
    plt.vlines(x_positions[0], line_y*0.95, line_y, color='black', lw=1)
    plt.vlines(x_positions[1], line_y*0.95, line_y, color='black', lw=1)

    plt.tight_layout()
    plt.savefig(os.path.join(box_save_path, f'{col}.pdf'))
    plt.close()

    # ================== 统计检验
    x0 = test_signature_df[test_signature_df['ALN status'] == 0][col]
    x1 = test_signature_df[test_signature_df['ALN status'] == 1][col]

    mw_stat, mw_p = mannwhitneyu(x0, x1, alternative='two-sided')
    ks_stat, ks_p = ks_2samp(x0, x1)

    results.append({
        'Feature': col,
        'Mann-Whitney U p-value': mw_p,
        'K-S Test p-value': ks_p
    })

# ======================================
# 汇总检验结果保存为CSV
results_df = pd.DataFrame(results)
results_df.to_csv(os.path.join(signature_analysis_result_path, 'Mann-Whitney_K-S.csv'), index=False)
print(results_df)




from scipy.stats import fisher_exact, chi2_contingency

# 你定义的特征列（每个是一个 signature）
signature_cols = [col for col in signature_X.columns]

results = []

for sig in signature_cols:
    # 中位数分组
    median_val = signature_X[sig].median()
    group = signature_X[sig] > median_val  # True 为 high，False 为 low

    # 生成混淆矩阵: 行为 high/low，列为 ALN=1/0
    group_high_y1 = ((group == True) & (signature_y == 1)).sum()
    group_high_y0 = ((group == True) & (signature_y == 0)).sum()
    group_low_y1 = ((group == False) & (signature_y == 1)).sum()
    group_low_y0 = ((group == False) & (signature_y == 0)).sum()

    contingency_table = [[group_high_y1, group_high_y0],
                         [group_low_y1, group_low_y0]]

    # 使用 Fisher's 精确检验或卡方检验（适用于计数较小的情况）
    if any(x < 5 for row in contingency_table for x in row):
        oddsratio, p = fisher_exact(contingency_table)
        test_type = "Fisher"
    else:
        chi2, p, _, _ = chi2_contingency(contingency_table)
        test_type = "Chi2"

    results.append({
        'signature': sig,
        'high ALN+ (n)': group_high_y1,
        'high ALN- (n)': group_high_y0,
        'low ALN+ (n)': group_low_y1,
        'low ALN- (n)': group_low_y0,
        'p_value': p,
        'test_type': test_type
    })

    plot_df = pd.DataFrame({
            'Group': ['High', 'High', 'Low', 'Low'],
            'ALN status': ['ALNM', 'Non-ALNM', 'ALNM', 'Non-ALNM'],
            'Count': [group_high_y1, group_high_y0, group_low_y1, group_low_y0]
        })

    # 转化为比例
    plot_df['Proportion'] = plot_df.groupby('Group')['Count'].transform(lambda x: x / x.sum())
    plot_dir = opj(signature_analysis_result_path,'barplot')
    md(plot_dir)
    # plt.figure(figsize=(6,4))

    # sns.barplot(data=plot_df, x='Group', y='Proportion', hue='ALN status', palette={'ALNM':'#FF6666','Non-ALNM':'#66B2FF'},width=0.2)


    # line_y = 0.8  # 横线高度
    # plt.plot([0,1], [line_y, line_y], color='black', lw=1)  # 横线
    # if p < 0.001:
    #     signif2 = '***'
    # elif mw_p < 0.01:
    #     signif2 = '**'
    # elif mw_p < 0.05:
    #     signif2 = '*'
    # else:
    #     signif2 = ''
    # plt.text(0.5, line_y * 1.12, f"p = {format_pval_scientific(p)}",  ha='center', va='bottom', fontsize=11)
    # plt.text(0.5, line_y * 1.01, f'{signif2}', ha='center', va='bottom', fontsize=9)

    # # 两端垂直短线
    # plt.vlines(0, line_y*0.985, line_y, color='black', lw=1)
    # plt.vlines(1, line_y*0.985, line_y, color='black', lw=1)
    # plt.legend(frameon=False)
    # plt.title(f'{sig} High/Low Signature ALN status Distribution')
    # plt.ylabel('Proportion')
    # plt.ylim(0,1)
    # plt.tight_layout()
    # plt.savefig(opj(plot_dir, f'{sig}_high_low_ALN_barplot.pdf'))
    # plt.close()
    groups = ['High', 'Low']
    x_positions = np.array([0.3, 0.6])
    width = 0.1  # 每根柱子的宽度

    # 两类 ALN 状态比例
    proportions_ALNM = plot_df[plot_df['ALN status']=='ALNM']['Proportion'].values
    proportions_NonALNM = plot_df[plot_df['ALN status']=='Non-ALNM']['Proportion'].values

    plt.figure(figsize=(5,4))

    # 绘制 ALNM
    plt.bar(x_positions - width/2, proportions_ALNM, width=width, color='#FF6666', label='ALNM')
    # 绘制 Non-ALNM
    plt.bar(x_positions + width/2, proportions_NonALNM, width=width, color='#66B2FF', label='Non-ALNM')

    plt.xlim(0,0.9)
    plt.xticks(x_positions, groups)
    plt.ylabel('Proportion')
    # plt.legend(frameon=False)
    plt.title(f'High/Low {sig} Signature vs ALN Status')
    plt.savefig(opj(plot_dir, f'{sig}_high_low_ALN_barplot.pdf'))
    plt.tight_layout()
    plt.show()

# ================== 保存统计结果
results_df = pd.DataFrame(results)
results_df = results_df.sort_values('p_value')
results_df.to_csv(opj(signature_analysis_result_path, 'signature_high_low_ALN_test_results.csv'), index=False)
